import fastapi
from fastapi import FastAPI, HTTPException
from langchain.vectorstores import FAISS
from langchain.embeddings import HuggingFaceEmbeddings
from pydantic import BaseModel, Field
from typing import List, Optional, Dict
import os
import pathlib
import time
import requests
import json
import argparse

flags = argparse.ArgumentParser(description="Knowledge Server API")
flags.add_argument("--port", type=int, default=8000, help="Port to run the server on")
flags.add_argument("--remote_host", type=str, default="localhost", help="AI Embedding server on")
flags.add_argument("--token", type=str, default="", help="API token for AI Embedding server")



class Document(BaseModel):
    id: Optional[str] = ""
    page_content: str
    metadata: Dict

class MessageRequest(BaseModel):
    query: Optional[str]
    name: str = "default"
    file: Optional[str]
    stream: Optional[bool] = False


VEC_PATH = pathlib.Path("~/.vecs").expanduser()
if not VEC_PATH.exists():
    VEC_PATH.mkdir(parents=True)

REMOTE_HOST = os.getenv("V_REMOTE_HOST", "localhost")
TOKEN = os.getenv("V_TOKEN", "")
PORT = int(os.getenv("V_PORT", "8000"))
class Embeding:

    def __init__(self, remote_host: str, token=""):
        self.remote_host = remote_host
        self.use_code = False
        self.token = token
        self.functions = None
    
    def embed_documents_remote(self, texts):
        uri = f"https://{self.remote_host}/v1/embeddings"
        # user_id = md5(time.asctime().encode()).hexdigest()
        try:
            m = {
                "messages":texts,
                "stream":True
            }
            ss = []
            for i in R(uri, m, use_stream=True, token=self.token):
                ss += i["data"]
            return ss
        except Exception as e:
            raise e
    
    def embed_query_remote(self, text):
        uri = f"https://{self.remote_host}/v1/embeddings"
        try:
            m = {
                "messages":[text],
                "stream":True
            }
            ss = []
            for i in R(uri, m, use_stream=True,token=self.token):
                ss += i["data"]
            return ss[0]
        except Exception as e:
            print(e)
            import ipdb;ipdb.set_trace()

    def embed_documents(self, texts):
        assert self.remote_host is not None
        return self.embed_documents_remote(texts)
        
    def embed_query(self, text: str):
        assert self.remote_host is not None
        return self.embed_query_remote(text)
        
    
    def __call__(self, texts):
        if isinstance(texts, str):
            return self.embed_query(texts)
        return self.embed_documents(texts)




def R(uri, object=None, method='post',use_stream=False,token="", **datas):
    data = {}
    data.update(datas)
    o = None
    if object is not None:
        o = object
        
        data.update(o)
    H = {
        "Content-Type": "application/json",
        "Authorization": "Bearer " + token,
    }
    if method == "post":
        response = requests.post(uri, headers=H, json=data, stream=use_stream, verify=False)
    else:
        response = requests.get(uri, json=data, stream=use_stream,verify=False)
    if response.status_code == 200:
        if use_stream:
            # 处理流式响应
            if "messages" in o:
                T = (len(o["messages"]) // 50) + 1
                if uri.endswith("embeddings"):
                    T = (len(o["messages"]) // 100) + 1
                # bar = tqdm.tqdm(total=T,desc=" + deal data")
                # bar.leave = False
                for line in response.iter_lines():
                    if line:
                        if line[:6] == b": ping" and line[6] != b"{":continue
                        decoded_line = line.decode('utf-8')[6:].strip()
                        if decoded_line:
                            try:
                                response_json = json.loads(decoded_line)
                                # bar.update(1)
                                yield response_json
                            except Exception as e:
                                print("Special Token:", e, len(line), line)
                # bar.clear()
                # bar.close()
        else:
            # 处理非流式响应
            decoded_line = response.json()
            yield decoded_line
    else:
        print("Error:", response.status_code)
        return None


def get_db(remote_host, name,docs=[], token="" ):
    embeddings = Embeding(remote_host=remote_host, token=token)
    vs_path = str(VEC_PATH /name )
    if os.path.isdir(vs_path) and os.path.exists(str(pathlib.Path(vs_path) / "index.faiss")):
        try:
            db = FAISS.load_local(vs_path, embeddings, allow_dangerous_deserialization=True)
        except Exception as e:
            db = FAISS.load_local(vs_path, embeddings)
        if len(docs) > 0:
            db.add_documents(docs)
    else:
        if len(docs) > 0:
            db = FAISS.from_documents(docs, embeddings)
        else:
            return None
    db.save_local(vs_path)
    return db

def search_vectorstore(query: str, vectorstore: FAISS, top_k=7):
    # 实现搜索功能
    docs = vectorstore.similarity_search(query, k=top_k)
    return docs

app = FastAPI()

@app.post("/search")
async def search(request: MessageRequest):
    # 实现搜索功能
    query = request.query
    name = request.name
    db = get_db(REMOTE_HOST, name, token=TOKEN)
    if db is not None:
        docs = search_vectorstore(query, db)
        return {"docs": docs}
    else:
        return {"docs": []}

@app.post("/create")
async def search(request: MessageRequest):
    # 实现搜索功能
    docs = []
    name = request.name
    file = request.file
    if file != "" and len(docs) == 0:
        with open(file, "r", encoding="utf-8") as f:
            docs.append(Document(id=name,page_content=f.read(), metadata={"source": file}))
    if len(docs) > 0:
        get_db(REMOTE_HOST, name,docs, token=TOKEN)
    return {"docs": "ok"}

import uvicorn

args = flags.parse_args()
PORT = args.port
REMOTE_HOST = args.remote_host
TOKEN = args.token
uvicorn.run(app,log_level="info", host='0.0.0.0', port=PORT, workers=1)